# -*- coding: utf-8 -*-
# 1.加载CSV文件中的数据，并反转数据以便按照时间顺序排列
import matplotlib.pyplot as plt
import numpy as np
import torch

torch.manual_seed(42)
data = np.loadtxt('data-02-stock_daily.csv', delimiter=',')
data = data[::-1]
# 2.对数据进行归一化处理
from sklearn.preprocessing import MinMaxScaler

data = MinMaxScaler().fit_transform(data)
# 3.设置时间步长
c = 7
# 4.初始化输入和输出数据列表
x = []
y = []
# 5.遍历数据，构建输入和输出
for i in range(len(data) - c):
    x.append(data[i:i + c])
    y.append(data[i + c, -1])
# 6.将输入和输出数据转换为Tensor
x = torch.tensor(x, dtype=torch.float)
y = torch.tensor(y, dtype=torch.float).reshape(-1, 1)
# 7.划分数据集为训练集和测试集
from sklearn.model_selection import train_test_split

x_train, x_test, y_train, y_test = train_test_split(x, y, shuffle=False, random_state=42)
# 8.打印训练数据的形状
print(x_train.shape, y_train.shape)
# 9.设置输入数据的特征维度
input_size = x_train.shape[2]


# 10.定义RNN模型
class Rnn(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.rnn1 = torch.nn.RNN(input_size=input_size, hidden_size=50, batch_first=True)
        self.rnn2 = torch.nn.RNN(input_size=50, hidden_size=50, batch_first=True)
        self.fc = torch.nn.Linear(in_features=50, out_features=1)

    def forward(self, x):
        x, _ = self.rnn1(x)
        x, _ = self.rnn2(x)
        out = self.fc(x[:, -1, :])
        return out


# 11.实例化模型
model = Rnn()
# 12.定义损失函数
loss_fn = torch.nn.MSELoss()
# 13.定义优化器
optim_adam = torch.optim.Adam(model.parameters())
# 14.初始化损失列表
loss_list = []
# 15.训练模型
model.train()
for i in range(2000):
    optim_adam.zero_grad()
    h = model(x_train)
    loss = loss_fn(h, y_train)
    loss.backward()
    optim_adam.step()
    if i % 100 == 0:
        print(i + 1, loss)
# 16.模型预测测试集
model.eval()
with torch.no_grad():
    h = model(x_test)
    loss = loss_fn(h, y_test)
    # 17.打印预测结果
    print(loss)
    # 18.绘制预测结果和真实结果
    plt.plot(h)
    plt.plot(y_test)
    plt.show()
